import pandas as pd

from .time_df import TimeDF, fit_nan


def lag(Timedf,n):
      df=Timedf.df
      df=df.shift(n)
      return TimeDF(df)

def count(Timedf,n):
      df=Timedf.df.copy()
      df=df.rolling(n).sum()
      nan_df=fit_nan(df)
      df=(df==n)+nan_df
      return TimeDF(df)

def _and(*args):
      res=args[0]
      for p in args:
            res=res&p
      return res

def _or(*args):
      res=args[0]
      for p in args:
            res=res|p
      return res

def eq(Timedf,value):
      return Timedf==value

def ge(Timedf,value):
      return Timedf>=value

def le(Timedf,value):
      return Timedf<=value

def gt(Timedf,value):
      return Timedf>value

def lt(Timedf,value):
      return Timedf<value

_FUNCTIONS={
      'lag':lag,
      '_and': _and,
      '_or': _or,
      'eq': eq,
      'ge': ge,
      'gt': gt,
      'le': le,
      'lt':lt,
      'count':count
}
